/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */

/**
 *
 * @author ASUS
 */

import java.util.*;
import java.util.concurrent.*;

public class GeneticAlgorithmTrainer implements TrainingAlgorithm {
    private int populationSize;
    private double mutationRate;
    private Random random = new Random();
    
    public GeneticAlgorithmTrainer(int populationSize, double mutationRate) {
        this.populationSize = populationSize;
        this.mutationRate = mutationRate;
    }
    
    @Override
    public NeuralNetwork train(double[][] inputs, double[] targets, int generations) {
        System.out.println("开始遗传算法训练...");
        
        // 创建初始种群
        List<GeneticAlgorithmNN> population = new ArrayList<>();
        for (int i = 0; i < populationSize; i++) {
            population.add(new GeneticAlgorithmNN());
        }
        
        GeneticAlgorithmNN bestIndividual = null;
        double bestLoss = Double.MAX_VALUE;
        
        // 记录训练历史
        List<Double> bestLossHistory = new ArrayList<>();
        
        for (int gen = 0; gen < generations; gen++) {
            // 评估种群
            List<IndividualFitness> fitnessList = evaluatePopulation(population, inputs, targets);
            
            // 找到最佳个体
            IndividualFitness currentBest = Collections.max(fitnessList, 
                Comparator.comparingDouble(f -> f.fitness));
            
            double currentBestLoss = 1.0 / currentBest.fitness - 1.0;
            
            if (currentBestLoss < bestLoss) {
                bestLoss = currentBestLoss;
                bestIndividual = currentBest.individual;
            }
            
            bestLossHistory.add(bestLoss);
            
            // 输出进度
            if (gen % 10 == 0) {
                System.out.printf("代数 %3d: 最佳损失 = %.6f%n", gen, bestLoss);
            }
            
            // 停止条件
            if (bestLoss < 0.01 || (gen > 50 && isConverged(bestLossHistory))) {
                System.out.printf("提前停止在代数 %d, 最终损失 = %.6f%n", gen, bestLoss);
                break;
            }
            
            // 创建新一代
            population = createNewGeneration(fitnessList);
        }
        
        System.out.printf("遗传算法训练完成，最终损失 = %.6f%n", bestLoss);
        return bestIndividual;
    }
    
    /**
     * 评估种群适应度 - 使用更敏感的适应度函数
     */
    private List<IndividualFitness> evaluatePopulation(List<GeneticAlgorithmNN> population, 
                                                     double[][] inputs, double[] targets) {
        List<IndividualFitness> fitnessList = new ArrayList<>();
        
        for (GeneticAlgorithmNN individual : population) {
            double loss = individual.calculateLoss(inputs, targets);
            // 使用指数适应度函数，对小损失更敏感
            double fitness = Math.exp(-10 * Math.sqrt(loss));
            fitnessList.add(new IndividualFitness(individual, fitness, loss));
        }
        
        return fitnessList;
    }
    
    /**
     * 创建新一代种群
     */
    private List<GeneticAlgorithmNN> createNewGeneration(List<IndividualFitness> fitnessList) {
        List<GeneticAlgorithmNN> newPopulation = new ArrayList<>();
        
        // 按适应度排序
        fitnessList.sort((f1, f2) -> Double.compare(f2.fitness, f1.fitness));
        
        // 精英选择：保留前30%
        int eliteCount = Math.max(2, populationSize / 3);
        for (int i = 0; i < eliteCount; i++) {
            newPopulation.add(fitnessList.get(i).individual);
        }
        
        // 生成剩余个体
        while (newPopulation.size() < populationSize) {
            // 70%概率使用交叉，30%概率使用变异
            if (random.nextDouble() < 0.7) {
                // 交叉：选择两个较好的父代
                GeneticAlgorithmNN parent1 = selectParent(fitnessList);
                GeneticAlgorithmNN parent2 = selectParent(fitnessList);
                GeneticAlgorithmNN child = aggressiveCrossover(parent1, parent2);
                newPopulation.add(child);
            } else {
                // 变异：选择一个父代进行变异
                GeneticAlgorithmNN parent = selectParent(fitnessList);
                GeneticAlgorithmNN child = aggressiveMutate(parent);
                newPopulation.add(child);
            }
        }
        
        return newPopulation;
    }
    
    /**
     * 选择父代 - 基于适应度的概率选择
     */
    private GeneticAlgorithmNN selectParent(List<IndividualFitness> fitnessList) {
        double totalFitness = fitnessList.stream().mapToDouble(f -> f.fitness).sum();
        double randomValue = random.nextDouble() * totalFitness;
        
        double cumulative = 0;
        for (IndividualFitness fitness : fitnessList) {
            cumulative += fitness.fitness;
            if (cumulative >= randomValue) {
                return fitness.individual;
            }
        }
        
        return fitnessList.get(0).individual;
    }
    
    /**
     * 激进的交叉操作
     */
    private GeneticAlgorithmNN aggressiveCrossover(GeneticAlgorithmNN parent1, GeneticAlgorithmNN parent2) {
        GeneticAlgorithmNN child = new GeneticAlgorithmNN();
        double[] params1 = parent1.getParameters();
        double[] params2 = parent2.getParameters();
        double[] childParams = new double[params1.length];
        
        // 多点交叉 + 随机扰动
        for (int i = 0; i < params1.length; i++) {
            if (random.nextDouble() < 0.5) {
                childParams[i] = params1[i];
            } else {
                childParams[i] = params2[i];
            }
            
            // 添加小扰动
            if (random.nextDouble() < 0.3) {
                childParams[i] += random.nextGaussian() * 0.2;
            }
        }
        
        child.setParameters(childParams);
        return child;
    }
    
    /**
     * 激进的变异操作
     */
    private GeneticAlgorithmNN aggressiveMutate(GeneticAlgorithmNN parent) {
        GeneticAlgorithmNN child = new GeneticAlgorithmNN();
        double[] parentParams = parent.getParameters();
        double[] childParams = parentParams.clone(); // 克隆父代参数
        
        // 变异多个参数
        int mutateCount = Math.max(1, (int)(childParams.length * mutationRate));
        for (int i = 0; i < mutateCount; i++) {
            int index = random.nextInt(childParams.length);
            // 较大的变异幅度
            childParams[index] += random.nextGaussian() * 0.5;
            // 限制参数范围
            childParams[index] = Math.max(-2, Math.min(2, childParams[index]));
        }
        
        child.setParameters(childParams);
        return child;
    }
    
    /**
     * 检查是否收敛
     */
    private boolean isConverged(List<Double> lossHistory) {
        if (lossHistory.size() < 20) return false;
        
        // 检查最近20代是否没有显著改进
        double recentImprovement = 0;
        for (int i = lossHistory.size() - 20; i < lossHistory.size() - 1; i++) {
            recentImprovement += lossHistory.get(i) - lossHistory.get(i + 1);
        }
        
        return recentImprovement < 0.001;
    }
    
    /**
     * 个体适应度记录
     */
    private static class IndividualFitness {
        final GeneticAlgorithmNN individual;
        final double fitness;
        final double loss;
        
        IndividualFitness(GeneticAlgorithmNN individual, double fitness, double loss) {
            this.individual = individual;
            this.fitness = fitness;
            this.loss = loss;
        }
    }
}